import asyncio
import uuid
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select

from app.services.model_router import call_model
from app.db.session import get_session
from app.db.redis import get_cached_response, set_cached_response
from app.models.message import Message
from app.models.session import Session
from app.models.history import History  # ✅ 新增导入

router = APIRouter()


@router.websocket("/chat")
async def websocket_chat(ws: WebSocket, db: AsyncSession = Depends(get_session)):
    await ws.accept()
    session_id = ws.query_params.get("session_id") or str(uuid.uuid4())

    # 检查 session 是否存在
    existing_session = await db.get(Session, session_id)
    if not existing_session:
        new_session = Session(id=session_id)
        db.add(new_session)
        await db.commit()

    try:
        while True:
            data = await ws.receive_json()
            model = data.get("model") or "qianfan"
            message = data.get("message") or "1+1=?"

            # ==========================
            # 🚀 Redis 缓存
            # ==========================
            cached = await get_cached_response(model, message)
            if cached:
                reply = cached
            else:
                reply = await call_model(model, message)
                await set_cached_response(model, message, reply)

            # ==========================
            # 💬 保存用户消息
            # ==========================
            user_msg = Message(session_id=session_id, sender="user", content=message)
            db.add(user_msg)
            await db.commit()

            # ==========================
            # 🧠 流式发送 AI 回复
            # ==========================
            for i in range(0, len(reply), 5):
                chunk = reply[i:i+5]
                await ws.send_text(chunk)
                await asyncio.sleep(0.2)

            # ==========================
            # 🤖 保存 AI 回复
            # ==========================
            ai_msg = Message(session_id=session_id, sender="ai", content=reply)
            db.add(ai_msg)

            # ==========================
            # 🧾 更新或创建历史记录
            # ==========================
            result = await db.execute(
                select(History).where(History.session_id == session_id)
            )
            history = result.scalars().first()

            if not history:
                new_history = History(
                    session_id=session_id,
                    model_name=model,  # ✅ 保存模型名
                    summary=reply[:100],  # 简要摘要
                    last_message=reply
                )
                db.add(new_history)
            else:
                history.model_name = model
                history.last_message = reply
                history.updated_at = None  # 触发自动更新时间
                db.add(history)

            await db.commit()

    except WebSocketDisconnect:
        print(f"❌ WebSocket {session_id} disconnected")
